import torch

import runner2
import dataset
import model2
import assistant
import tokenizer
from ai.config.config import GxlNode
from project.rr_shiyou_zuoye.my_neuro_model import runner2_1, model1

if __name__ == "__main__":
    charset_normalizer = tokenizer.ChineseCharacterTokenizer()
    charset_normalizer.load_vocab()
    model = model1.LinearModel()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-5)
    loss_fun = torch.nn.CrossEntropyLoss()
    lr_schedule = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
    train_iter, dev_iter, test_iter = dataset.get_dataloader('../my_spider_douban_formal/output/clean_data.jsonl', 16)
    config = GxlNode.get_config_from_yaml('./config.yaml')
    runner_man = runner2_1.RunnerGxl(model, optimizer, loss_fun, train_iter, config, multi=False,
                                     device=torch.device('cpu'), scheduler=lr_schedule, valid_loader=dev_iter)
    runner_man.run(1000)
